import torch
import torch.nn.functional as F
import numpy as np
import os
import cv2
import pandas as pd

from tqdm import tqdm
from skimage.measure import label, regionprops

import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')

def mask2rle(mask):
    ''' takes a 2d boolean numpy array and turns it into a space-delimited RLE string '''
    mask = mask.T.reshape(-1) # make 1D, column-first
    mask = np.pad(mask, 1, mode="constant") # make sure that the 1d mask starts and ends with a 0
    starts = np.nonzero((~mask[:-1]) & mask[1:])[0] # start points
    ends = np.nonzero(mask[:-1] & (~mask[1:]))[0] # end points
    rle = np.empty(2 * starts.size, dtype=int) # interlacing...
    rle[0::2] = starts + 1# ...starts...
    rle[1::2] = ends - starts # ...and lengths
    rle = ' '.join([ str(elem) for elem in rle ]) # turn into space-separated string
    return rle

out_info = {}
score_threshold = 0.4
model_flags = ["UNet_Resnest02", "UNet_Resnest03"]
dataroot = "E:/睡眠分期数据/hubmap-kidney-segmentation/test/"
filelist = list(filter(lambda fn:fn.endswith(".tiff"), os.listdir(dataroot)))

main_bar = tqdm(filelist)
for fn in main_bar:
    pred_mask = None
    for mf in model_flags:
        if pred_mask is None:
            pred_mask = np.load("./submission/{}/{}_rawpred.npy".format(mf, fn))
        else:
            pred_mask += np.load("./submission/{}/{}_rawpred.npy".format(mf, fn))
    pred_mask = pred_mask/len(model_flags)
    pred_mask = (pred_mask >= score_threshold).astype(np.uint8)
    out_info[len(out_info)] = {'id':fn.split(".")[0], 'predicted': mask2rle(pred_mask)}

submission = pd.DataFrame.from_dict(out_info, orient='index')
submission.to_csv('./submission/submission.csv', index=False)